import promptToString from '@/common/utils/promptToString';
import {
  LORATRIGGER_PROMPT_SEPARATOR,
  SECONDARY_PROMPT_SEPARATOR,
} from '@/config/constants';

export const getPromptLora = (prompt: string) => {
  const loraRegex = /withLora\(([^\\)]+)\)/g; // Scan For Lora Syntax
  const loraMatches = [...prompt.matchAll(loraRegex)]; // Match All Lora Syntaxes

  const loraMatchNames: string[] = [];
  if (loraMatches.length > 0) {
    loraMatches.forEach((loraMatch) => {
      const loraName = loraMatch[1].split(',')[0];
      loraMatchNames.indexOf(loraName) === -1 && loraMatchNames.push(loraName); // Add Loras In Prompt
    });
  }
  return { loraMatches, loraMatchNames };
};
export function getPromptAndNegative(inputPrompt: any) {
  let prompt: string =
    typeof inputPrompt === 'string' ? inputPrompt : promptToString(inputPrompt);

  let negativePrompt = '';

  // Matches all negative prompts, 1st capturing group is the prompt itself
  const negativePromptRegExp = new RegExp(/\[([^\][]*)]/, 'gi');

  // Grab the actual prompt matches (capturing group 1 is 1st index of match)
  const negativePromptMatches = [...prompt.matchAll(negativePromptRegExp)].map(
    (match) => match[1]
  );

  if (negativePromptMatches.length) {
    // Build the negative prompt itself
    negativePrompt = negativePromptMatches.join(' ');

    // Replace each match, including its surrounding brackets
    // Remove each pair of empty brackets
    // Trim whitespace
    negativePromptMatches.forEach((match) => {
      prompt = prompt.replace(`[${match}]`, '').replaceAll('[]', '').trim();
    });
  }
  // 去除lora
  const { loraMatchNames } = getPromptLora(prompt);

  if (loraMatchNames.length) {
    loraMatchNames.forEach((loraName) => {
      const loraRegex = new RegExp(
        `withLora\\(${loraName},?\\s*([^\\)]+)?\\)`,
        'g'
      );
      prompt = prompt.replaceAll(loraRegex, '').trim();
    });
  }

  // 去除辅助提示词
  prompt = prompt.split(SECONDARY_PROMPT_SEPARATOR)[0];

  // 去除lora触发词
  let loraTriggerTokensString = '';
  if (
    loraMatchNames.length &&
    prompt.indexOf(LORATRIGGER_PROMPT_SEPARATOR) >= 0
  ) {
    const [loraTriggerTokens, _p] = prompt.split(LORATRIGGER_PROMPT_SEPARATOR);
    if (_p) {
      prompt = _p;
      loraTriggerTokensString = loraTriggerTokens;
    } else {
      prompt = loraTriggerTokens;
    }
  }

  return [
    prompt,
    negativePrompt,
    loraMatchNames.join(),
    loraTriggerTokensString,
  ];
}
